/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.mcp.sample.client;

import java.util.Map;

import io.modelcontextprotocol.client.McpClient;
import io.modelcontextprotocol.spec.McpClientTransport;
import io.modelcontextprotocol.spec.McpSchema.CallToolRequest;
import io.modelcontextprotocol.spec.McpSchema.CallToolResult;
import io.modelcontextprotocol.spec.McpSchema.ListToolsResult;

/**
 * @author Christian Tzolov
 */

public class SampleClient {

	private final McpClientTransport transport;

	public SampleClient(McpClientTransport transport) {
		this.transport = transport;
	}

	public void run() {

		var client = McpClient.sync(this.transport).build();

		client.initialize();

		client.ping();

		// List and demonstrate tools
		ListToolsResult toolsList = client.listTools();
		System.out.println("Available Tools = " + toolsList);
		System.out.println();

		// CallToolResult weatherForcastResult = client.callTool(new CallToolRequest("getWeatherForecastByLocation",
		// 		Map.of("latitude", "47.6062", "longitude", "-122.3321")));
		// System.out.println("Weather Forcast: " + weatherForcastResult);
		//
		// CallToolResult alertResult = client.callTool(new CallToolRequest("getAlerts", Map.of("state", "NY")));
		// System.out.println("Alert Response = " + alertResult);


		// 测试查看表行数
		CallToolResult getTableRowCountResult = client.callTool(new CallToolRequest("getTableRowCount",
				Map.of("tableName", "pc_file")));
		System.out.println("\nGet Table Row Count: " + getTableRowCountResult);

		// 测试创建数据源
		CallToolResult createDatasourceResult = client.callTool(new CallToolRequest("createDatasource",
				Map.of("key", "internal", "url", "jdbc:mysql://localhost:3306/ydzb", "username", "root",
						"password", "123456", "databaseType", "mysql")));
		System.out.println("\nCreate Datasource: " + createDatasourceResult);

		CallToolResult showTablesResult = client.callTool(new CallToolRequest("listTables", Map.of("", "")));
		System.out.println("\nShow Tables: " + showTablesResult);


		// 测试查看数据源
		CallToolResult datasourceResult = client.callTool(new CallToolRequest("listDatasources", Map.of("", "")));
		System.out.println("\nDatasource list: " + datasourceResult);

		// 测试查看表行数
		CallToolResult getTableRowCountResult2 = client.callTool(new CallToolRequest("getTableRowCount",
				Map.of("tableName", "pc_file")));
		System.out.println("\nGet Table Row Count: " + getTableRowCountResult2);

		CallToolResult querySQLResult = client.callTool(new CallToolRequest("querySQL", Map.of("sql", "select * from pc_file limit 10")));
		System.out.println("\nQuery SQL: " + querySQLResult);

		// 测试创建数据源
		CallToolResult createDatasourceResult2 = client.callTool(new CallToolRequest("createDatasource",
				Map.of("key", "external", "url", "jdbc:mysql://localhost:3306/ydzb_external", "username", "root",
						"password", "123456", "databaseType", "mysql")));
		System.out.println("\nCreate Datasource: " + createDatasourceResult2);

		// 测试查看表行数
		CallToolResult getTableRowCountResult3 = client.callTool(new CallToolRequest("getTableRowCount",
				Map.of("tableName", "pc_file")));
		System.out.println("\nGet Table Row Count: " + getTableRowCountResult3);



		// 测试查看数据源
		CallToolResult datasourceResult1 = client.callTool(new CallToolRequest("listDatasources", Map.of("", "")));
		System.out.println("\nDatasource list: " + datasourceResult1);

		CallToolResult chooseDatabaseResult = client.callTool(new CallToolRequest("chooseDatabase", Map.of("key", "internal")));
		System.out.println("\nChoose Database: " + chooseDatabaseResult);

		// 测试查看表行数
		CallToolResult getTableRowCountResult4 = client.callTool(new CallToolRequest("getTableRowCount",
				Map.of("tableName", "pc_file")));
		System.out.println("\nGet Table Row Count: " + getTableRowCountResult4);

		client.closeGracefully();
	}

}
